import pickle
# import tarfile
import tensorflow as tf
import numpy as np


def unpickle_cifar_dic(file):
    """
    Helper function: unpickles a dictionary (used for loading CIFAR)
    :param file: filename of the pickle
    :return: tuple of (images, labels)
    """
    fo = open(file, 'rb')
    dict = pickle.load(fo, encoding='bytes')
    fo.close()
    d_decoded = {}
    for key, value in dict.items():
        d_decoded[key.decode('utf8')] = value
    return d_decoded["data"], d_decoded["labels"]


def load():
    train_files = ["data_batch_" + str(i) for i in range(1, 6)]

    test_file = ["test_batch"]
    # cifar10_files = train_files + test_file
    data_dir = "D:/code/pythonCode/PracticeTensorflow/data"

    images = []
    labels = []
    for file in train_files:
        filename = data_dir + "/cifar-10-batches-py/" + file

        images_tmp, labels_tmp = unpickle_cifar_dic(filename)

        images.append(images_tmp)
        labels.append(labels_tmp)

    train_data = np.asarray(images, dtype=np.float32).reshape((50000, 3, 32, 32))
    train_data = np.swapaxes(train_data, 1, 3)
    train_labels = np.asarray(labels, dtype=np.int32).reshape(50000)

    return train_data, train_labels


if __name__ == "__main__":
    data, label = load()
    # print((data, label))
    print(data.shape)
    print(label.shape)
